# This file contains an implementation of the sum-product decoding
# algorithm for LDPC codes.

# global variables storing variable node names and check node names
vNodes := {};
cNodes := {};


# This function builds a graph
# the number of a variable vertex is the corresponding column number
# of H. the number of a check vertex is the row number + total number
# of columns of H. The function returns a record with three fields.
# 1. Variable nodes;
# 2. Check nodes;
# 3. The constructed Tanner Graph.
TannerGraph := proc(H)
local rownum, colnum, E, i, j;
global vNodes, cNodes;
    (rownum, colnum) := LinearAlgebra:-Dimension(H);
    E := {};
    vNodes := {};
    cNodes := {};
    for i from 1 to rownum do
        for j from 1 to colnum do
            if H(i, j) = 1 then
                vNodes := vNodes union {j};
                cNodes := cNodes union {colnum + i};
                E := E union {{j, colnum + i}};
            end if;
        od;
    od;
    GraphTheory:-Graph(E)
end proc;

# a testing case for building the Tanner graph
#L := TannerGraph(GH[1]):
#GraphTheory:-DrawGraph(L[3]);

# This function computes the initial value for q_ij(0)
# which is 1 - Pr_i := 1 - Pr_i(ci = 1 | yi)
# given the channel error rate P. We assume the probability of
# 0 -> 1 and 1 -> 0 are both equivalent to P.
initialProb := proc(P, i) option inline;
    1 - P[i];
end proc;

# compute response from check node c_j.
# i.e., r_ji(0)
computeCheckNodeResponse2 := proc(qmat, i, j, neighbors) option inline;
    0.5 * (mul(2*qmat(k, j) -1, k = (convert(neighbors, set) minus {i})) + 1)
end proc;

computeCheckNodeResponse := proc(qmat, i, j, neighbors, multresult)
local a;
    a := 2*qmat(i, j) -1;
    if a <> 0 then
	return evalf(0.5 * (multresult/a) + 0.5);
    end if;	
    0.5 * (mul(2*qmat(k, j) -1, k = (convert(neighbors, set) minus {i})) + 1)
end proc;

# compute part of the response from variable node v_i.
multiplyResponse1 := proc(rmat, i, j,  neighbors, vSize, pi) option inline;
    (1-pi) * mul(rmat(k-vSize, i), k=(convert(neighbors, set) minus {(j+vSize)}));
end proc;

multiplyResponse2 := proc(rmat, i, j,  neighbors, vSize, pi) option inline;
    pi * mul(1 - rmat(k-vSize, i), k=(convert(neighbors, set) minus {(j+vSize)}));
end proc;

# check whether C is a codeword
isCodeWord := proc(H, C, cSize, vSize)
local i, j, sum;
    for i from 1 to cSize do
        sum := add(H(i, j) * C[j], j = 1..vSize);
        if (sum mod 2) = 1 then
            return false;
        end if;
    od;
    true
end proc;

# initialization of q_ij(0), q_ij(1)s are implied by 1 - q_ij(0).
initVariableNode := proc(qmat, vSize, P, tGraph)
local v, i, j, c, qij;
global vNodes;
    for v in vNodes do
        i := v;
        qij := 1 - P[i];
        for c in GraphTheory:-Neighbors(tGraph, v) do
            j := c - vSize;
            qmat(i, j) := qij;
        od;
    od;
end proc;

# compute updates to qmat(i, j) for the last step in decoding
computeVariableNodeResponse := proc(rmat, i, j, neighbors, vSize, pi, Q0i, Q1i)
local rmatji, coef0, coef1;
    if (pi = 1.) or (pi = 0.) then
        return  (1. - pi);
    end if;
#    elif (pi = 0.) then
#        return  1;
#    else
        coef0 := Q0i;
        coef1 := Q1i;
        rmatji := rmat(j, i);
        if rmatji = 0 then
            coef0 := multiplyResponse1(rmat, i, j,  neighbors, vSize, pi);
        elif rmatji = 1 then
            coef1 := multiplyResponse2(rmat, i, j,  neighbors, vSize, pi);
        else
            coef0 := coef0/rmatji;
            coef1 := coef1/(1-rmatji);
        end if;
        # return  
        evalf((1/(coef0 + coef1)) * coef0);
#    end if;
end proc;


# Global variables used in decodeSPA.
# We only allocate the following arrays and matrices once during
# the whole simulation for optimizing storage space.
# initialization of the vector containing estimated codewords that
# is transmitted
C  := Array(1..N);
Q0 := Array(1..N);
Q1 := Array(1..N);

# element q_ij is the message sent from variable node v_i to check
# node c_j.
qmat := Matrix(N, M, 0);

# element r_ji is the message sent from check node c_i to variable
# node v_i.
rmat := Matrix(M, N, 0);


# This function implements the LDPC decoder using the sum product
# algorithm in probability domain.
# S: received codeword bits.
# H: parity-check matrix .
# P: a list of error/erasure probabilities.
# tGraph: the tanner graph created from the parity check matrix.
decodeSPA := proc(S, H, lim, P, tGraph)
# global variables for saving storage space
global vNodes, cNodes, N, M;
global qmat, rmat, Q0, Q1; #qmat and rmat's elements are q_ij(0) and r_ij(0),
                     # Q stores probability pair [Q[0]/K, Q[1]/K]
                     #respectively;
global C; # stores the estimation of the transmitted codewords
          # determined at step 4.
### local variables
local neighbors; # the set of neighbors of a node.
local qij0, qij1, Kprod1, rmatji; # two symbolic expressions for solving Kij;
local coef0, coef1, prod0, prod1, HC;
local vSize, cSize, i, j, v, c, qij, rji, iter;
local multresult;

#    (vNodes, cNodes) := (triple[1], triple[2]);
    (vSize, cSize)   := (N, M);


    # 1. initialization of q_ij(0)
    # q_ij(1)s are implied by 1 - q_ij(0).
    initVariableNode(qmat, vSize, P, tGraph);

    iter := 0;
    while (iter < lim) do
        myprint("current iteration: ", iter);

        # 2. computes r_ji(0), where check node c_j and variable node v_i
        # are neighbors.
        for c in cNodes do
            j := c - vSize;
            neighbors := GraphTheory:-Neighbors(tGraph, c);
	    multresult := mul(2*qmat(k, j) -1, k = neighbors);
            for i in neighbors do
                rmat(j, i) := computeCheckNodeResponse(qmat, i, j, neighbors, multresult);
            od;
        od;

        # 3. compute the estimation of codewords.
        for v in vNodes do
            i := v;
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            prod0 := 1 - P[i];
            prod1 := 1 - prod0;
            if (prod0 = 0.) then
                prod1 := mul(1-rmat(c-vSize, i), c = neighbors);
            elif (prod1 = 0.) then
                prod0 := mul(rmat(c-vSize, i), c = neighbors);
            else
                for c in neighbors do
                    j := c - vSize;
                    prod0 := prod0 * rmat(j, i);
                    prod1 := prod1 * (1 - rmat(j, i));
                od;
            end if;
            Q0[i] := prod0;
            Q1[i] := prod1;
            Kprod1 := evalf(prod1/(prod0 + prod1));
            if Kprod1 > 0.5 then
                C[i] := 1;
            elif Kprod1 <= 0.5 then
                C[i] := 0;
            else
                print(Kprod1);
                print("overflow may happen!");
                error "terminate due to error";
            end if;
        od;

        if isCodeWord(H, C, cSize, vSize) then
             print("decoding is successful!");
             return [true, iter+1, convert(C, list)];
        end if;

        # 4. variable nodes update the response messages q_ij to check
        # nodes via equation solving.
        for v in vNodes do
            i := v;
            neighbors := GraphTheory:-Neighbors(tGraph, v);
            for c in neighbors do
                j := c - vSize;
                qmat(i, j) := computeVariableNodeResponse(rmat, i, j, neighbors, vSize, P[i], Q0[i], Q1[i]);
            od;
        od;
        iter := iter + 1;
    od;
    print("Warning: decoding terminated due to iteration limit!");
    [false, iter, convert(C, list)]
end proc;

